#include <iostream>
#include <opencv2/opencv.hpp>
#include <string>
#include <vector>

#include "common.hpp"
#include "det/postprocess.h"

int main()
{
    std::cout << "检测功能测试程序" << std::endl;

    // 创建测试图像 - 添加一些明显的目标
    cv::Mat test_image = cv::Mat::zeros(480, 640, CV_8UC3);

    // 添加一些矩形作为测试目标
    cv::rectangle(test_image, cv::Point(100, 100), cv::Point(300, 300), cv::Scalar(255, 255, 255), -1);
    cv::rectangle(test_image, cv::Point(400, 200), cv::Point(550, 350), cv::Scalar(128, 128, 128), -1);
    cv::rectangle(test_image, cv::Point(50, 350), cv::Point(200, 450), cv::Scalar(64, 64, 64), -1);

    // 添加文字说明
    cv::putText(test_image, "Test Targets", cv::Point(50, 50), cv::FONT_HERSHEY_SIMPLEX, 1.0, cv::Scalar(255, 255, 255),
                2);

    std::cout << "测试图像尺寸: " << test_image.cols << "x" << test_image.rows << std::endl;

    // 模型参数
    int model_in_h = 480;         // 高度
    int model_in_w = 640;         // 宽度
    float conf_threshold = 0.01f; // 很低的置信度阈值
    float nms_threshold = 0.45f;

    // 生成模拟的模型输出数据
    std::vector<uint16_t> output_data;
    int num_boxes = 6300;
    int box_size = 19;
    output_data.resize(num_boxes * box_size, 0);

    // 在中心位置添加一个高置信度的检测框
    int center_box_idx = 3150; // 中心位置
    uint16_t *box_ptr = output_data.data() + center_box_idx * box_size;

    // 设置边界框坐标 (归一化到0-255)
    box_ptr[0] = 127; // x center (0.5 * 255)
    box_ptr[1] = 127; // y center (0.5 * 255)
    box_ptr[2] = 100; // width (0.4 * 255)
    box_ptr[3] = 100; // height (0.4 * 255)
    box_ptr[4] = 200; // class confidence (0.8 * 255)

    std::cout << "设置检测框 " << center_box_idx << ": x=" << (int)box_ptr[0] << ", y=" << (int)box_ptr[1]
              << ", w=" << (int)box_ptr[2] << ", h=" << (int)box_ptr[3] << ", conf=" << (int)box_ptr[4] << std::endl;

    // 在另一个位置添加检测框
    int box2_idx = 1000;
    uint16_t *box2_ptr = output_data.data() + box2_idx * box_size;
    box2_ptr[0] = 200; // x center
    box2_ptr[1] = 150; // y center
    box2_ptr[2] = 80;  // width
    box2_ptr[3] = 80;  // height
    box2_ptr[4] = 150; // class confidence

    std::cout << "设置检测框 " << box2_idx << ": x=" << (int)box2_ptr[0] << ", y=" << (int)box2_ptr[1]
              << ", w=" << (int)box2_ptr[2] << ", h=" << (int)box2_ptr[3] << ", conf=" << (int)box2_ptr[4] << std::endl;

    std::cout << "生成测试输出数据，包含 " << num_boxes << " 个检测框" << std::endl;

    // 后处理参数
    BOX_RECT pads = {0, 0, 0, 0};
    float scale_w = (float)test_image.cols / model_in_w;
    float scale_h = (float)test_image.rows / model_in_h;

    std::cout << "缩放比例: scale_w=" << scale_w << ", scale_h=" << scale_h << std::endl;

    // 量化参数
    std::vector<int32_t> qnt_zps = {0};
    std::vector<float> qnt_scales = {1.0f};

    // 执行后处理
    DetectionResultsGroup results;
    int ret = post_process_yolov8(output_data.data(), model_in_h, model_in_w, conf_threshold, nms_threshold, pads,
                                  scale_w, scale_h, qnt_zps, qnt_scales, &results);

    if (ret == 0)
    {
        std::cout << "后处理成功，检测到 " << results.dets.size() << " 个目标" << std::endl;

        // 在图像上绘制检测结果
        for (const auto &det : results.dets)
        {
            std::cout << "检测到目标: " << det.det_name << " 置信度: " << det.score << " 位置: (" << det.box.x << ", "
                      << det.box.y << ", " << det.box.width << ", " << det.box.height << ")" << std::endl;

            // 绘制边界框
            cv::rectangle(test_image, det.box, cv::Scalar(0, 255, 0), 2);

            // 绘制标签
            std::string label = det.det_name + " " + std::to_string(det.score).substr(0, 4);
            cv::putText(test_image, label, cv::Point(det.box.x, det.box.y - 10), cv::FONT_HERSHEY_SIMPLEX, 0.5,
                        cv::Scalar(0, 255, 0), 2);
        }

        // 显示结果
        cv::imshow("Detection Test Results", test_image);
        cv::waitKey(0);
    }
    else
    {
        std::cout << "后处理失败，错误代码: " << ret << std::endl;
    }

    return 0;
}